package de.yvert.jingle.impl.openexr;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.zip.Deflater;

import de.yvert.jingle.hdr.HdrImage2D;
import de.yvert.jingle.help.Buffer;
import de.yvert.jingle.help.ConverterOutputStream;

public class ExrWriter
{

private static final int MAGIC = 20000630; // 30.06.2000

//private static final int FLAG_TILES = 2;

private ExrOutputStream out;
private ArrayList<Attribute> attributes = new ArrayList<Attribute>();

public ExrWriter(OutputStream out)
{
	this.out = new ExrOutputStream(new ConverterOutputStream(ByteOrder.LITTLE_ENDIAN, out));
}

public void write(HdrImage2D image) throws IOException
{
	int width = image.getWidth();
	int height = image.getHeight();
	
	Channel[] channels = new Channel[3];
	channels[0] = new Channel("B", PixelType.HALF);
	channels[1] = new Channel("G", PixelType.HALF);
	channels[2] = new Channel("R", PixelType.HALF);
	ChannelList channelList = new ChannelList(channels);
	
	CompressionType compression = CompressionType.ZIP_COMPRESSION;
	
	attributes.add(new ChannelListAttribute(StdAttributes.CHANNELS, channelList));
	attributes.add(new CompressionAttribute(StdAttributes.COMPRESSION, compression));
	attributes.add(new Box2iAttribute(StdAttributes.DATA_WINDOW, 0, 0, width-1, height-1));
	attributes.add(new Box2iAttribute(StdAttributes.DISPLAY_WINDOW, 0, 0, width-1, height-1));
	attributes.add(new LineOrderAttribute(StdAttributes.LINE_ORDER, LineOrder.INCREASING_Y));
	attributes.add(new FloatAttribute(StdAttributes.PIXEL_ASPECT_RATIO, 1));
	attributes.add(new Vector2fAttribute(StdAttributes.SCREEN_WINDOW_CENTER, 0, 0));
	attributes.add(new FloatAttribute(StdAttributes.SCREEN_WINDOW_WIDTH, 100));
//	attributes.add(new TileDescAttribute(StdAttributes.TILES, width, height, (byte) 0));
	
	out.writeInt(MAGIC);
	int vflags = 2; //(FLAG_TILES << 8) | 2;
	out.writeInt(vflags);
	for (Attribute attr : attributes)
		attr.write(out);
	out.writeByte((byte) 0);
	
	int blockSize = compression.getBlockSize();
	int numBlocks = (height+blockSize-1)/blockSize;
	long[] blockOffsets = new long[numBlocks];
	ByteArrayOutputStream tempOut = new ByteArrayOutputStream();
	ConverterOutputStream convOut = new ConverterOutputStream(ByteOrder.LITTLE_ENDIAN, tempOut);
	
	int maxBlockSize = blockSize*channelList.getLineLength(width);
	byte[] compressed = new byte[maxBlockSize];
	byte[] decompressed = new byte[maxBlockSize];
	byte[] pixeldata = new byte[maxBlockSize];
	Buffer buffer = new Buffer(Buffer.LITTLE_ENDIAN, pixeldata);
	
	float[] rgb = new float[3];
	
	for (int i = 0; i < numBlocks; i++)
	{
		blockOffsets[i] = tempOut.size();
		
		int ystart = i*blockSize;
		int ystop = ystart+blockSize;
		ystop = ystop < height ? ystop : height;
		int dataSize = (ystop-ystart)*channelList.getLineLength(width);
		
		// encode image to pixel data
		for (int y = ystart; y < ystop; y++)
		{
			int chsize = width*2;
			int offset = (y-ystart)*channelList.getLineLength(width);
			for (int x = 0; x < width; x++)
			{
				image.getRGB(x, y, rgb);
				buffer.setHalf(offset+0*chsize+2*x, rgb[2]);
				buffer.setHalf(offset+1*chsize+2*x, rgb[1]);
				buffer.setHalf(offset+2*chsize+2*x, rgb[0]);
			}
		}
		
		// reorder the pixel data
		for (int j = 0; j < dataSize/2; j++)
		{
			decompressed[j] = pixeldata[2*j+0];
			decompressed[j+dataSize/2] = pixeldata[2*j+1];
		}
		
		// predictor
		int p = decompressed[0] & 0xff;
		for (int j = 1; j < dataSize; j++)
		{
			int d = ((decompressed[j] & 0xff)-p+128) & 0xff;
			p = decompressed[j] & 0xff;
			decompressed[j] = (byte) d;
		}
		
		// compress
		Deflater compressor = new Deflater();
		compressor.setInput(decompressed, 0, dataSize);
		compressor.finish();
		int resultLength = compressor.deflate(compressed);
		compressor.end();
		
//		System.out.println(dataSize+" "+resultLength);
		
		convOut.writeInt(ystart);
		convOut.writeInt(resultLength);
		convOut.write(compressed, 0, resultLength);
	}
	
	for (int i = 0; i < numBlocks; i++)
		out.writeLong(blockOffsets[i]);
	out.write(tempOut.toByteArray());
	
	out.flush();
	System.out.println("finished writing exr");
}

}
